#include "EnginePrivate.h"

#define SHADER_CACHE_SIZE 1024

ShaderCompileManager GShaderCompileManager;

ShaderCompileManager::ShaderCompileManager()
{
	m_mapVSCache.clear();
	m_mapPSCache.clear();
	m_mapVSCodeBlob.clear();
	m_mapPSCodeBlob.clear();
}

ShaderCompileManager::~ShaderCompileManager()
{
	for(map<unsigned int, ID3D11VertexShader*>::iterator it = m_mapVSCache.begin(); it != m_mapVSCache.end(); ++it)
	{
		it->second->Release();
	}
	m_mapVSCache.clear();

	for(map<unsigned int, ID3D11PixelShader*>::iterator it = m_mapPSCache.begin(); it != m_mapPSCache.end(); ++it)
	{
		it->second->Release();
	}
	m_mapPSCache.clear();

	for(map<unsigned int, ID3D10Blob*>::iterator it = m_mapVSCodeBlob.begin(); it != m_mapVSCodeBlob.end(); ++it)
	{
		it->second->Release();
	}
	m_mapVSCodeBlob.clear();

	for(map<unsigned int, ID3D10Blob*>::iterator it = m_mapPSCodeBlob.begin(); it != m_mapPSCodeBlob.end(); ++it)
	{
		it->second->Release();
	}
	m_mapPSCodeBlob.clear();
}


ID3D11VertexShader* ShaderCompileManager::CompileVSShader(const char* pShaderFileName, const char* pShaderMainFun)
{
	// Compute Shader Cache Key
	unsigned int ShaderCacheKey = _ComputeShaderCacheKey(pShaderFileName, pShaderMainFun);

	ID3D11VertexShader* pShader = _FindVSShaderInCache(ShaderCacheKey);

	if(!pShader)
	{
		ID3D11Device* pGD3D11Device = gEngine.GetD3DDevice();

		ID3D10Blob* pCodeBlob = 0;
		ID3D10Blob* pErrorMsgBlob = 0;

		string szShaderString = _ReadShaderToString(pShaderFileName);

		HRESULT hResult = D3DCompile(szShaderString.c_str(), szShaderString.size(), pShaderFileName, NULL, NULL, pShaderMainFun, "vs_5_0", D3D10_SHADER_OPTIMIZATION_LEVEL3, 0, &pCodeBlob, &pErrorMsgBlob);

		assert(hResult == S_OK);

		// Process Compile error or warning message
		if(pErrorMsgBlob)
		{
			char* pErrorMsgString = (char*)pErrorMsgBlob->GetBufferPointer();
			pErrorMsgBlob->Release();
			pErrorMsgBlob = 0;
		}

		if(hResult != S_OK)
		{
			return 0;
		}

		assert(pCodeBlob);

		hResult = pGD3D11Device->CreateVertexShader(pCodeBlob->GetBufferPointer(), pCodeBlob->GetBufferSize(), 0, &pShader);

		assert(hResult == S_OK);

		m_mapVSCache.insert(pair<unsigned int, ID3D11VertexShader*>(ShaderCacheKey, pShader));

		m_mapVSCodeBlob.insert(pair<unsigned int, ID3D10Blob*>(ShaderCacheKey, pCodeBlob));
	}

	return pShader;
}

ID3D11PixelShader* ShaderCompileManager::CompilePSShader(const char* pShaderFileName, const char* pShaderMainFun)
{
	// Compute Shader Cache Key
	unsigned int ShaderCacheKey = _ComputeShaderCacheKey(pShaderFileName, pShaderMainFun);

	ID3D11PixelShader* pShader = _FindPSShaderInCache(ShaderCacheKey);

	if(!pShader)
	{
		ID3D11Device* pGD3D11Device = gEngine.GetD3DDevice();

		ID3D10Blob* pCodeBlob = 0;
		ID3D10Blob* pErrorMsgBlob = 0;

		string szShaderString = _ReadShaderToString(pShaderFileName);

		HRESULT hResult = D3DCompile(szShaderString.c_str(), szShaderString.size(), pShaderFileName, NULL, NULL, pShaderMainFun, "ps_5_0", D3D10_SHADER_OPTIMIZATION_LEVEL3, 0, &pCodeBlob, &pErrorMsgBlob);

		assert(hResult == S_OK);

		// Process Compile error or warning message
		if(pErrorMsgBlob)
		{
			char* pErrorMsgString = (char*)pErrorMsgBlob->GetBufferPointer();
			pErrorMsgBlob->Release();
			pErrorMsgBlob = 0;
		}

		if(hResult != S_OK)
		{
			return 0;
		}

		assert(pCodeBlob);

		hResult = pGD3D11Device->CreatePixelShader(pCodeBlob->GetBufferPointer(), pCodeBlob->GetBufferSize(), 0, &pShader);

		assert(hResult == S_OK);

		m_mapPSCache.insert(pair<unsigned int, ID3D11PixelShader*>(ShaderCacheKey, pShader));

		m_mapPSCodeBlob.insert(pair<unsigned int, ID3D10Blob*>(ShaderCacheKey, pCodeBlob));
	}

	return pShader;
}

ID3D10Blob* ShaderCompileManager::FindVSCodeBlob(const char* pShaderFileName, const char* pShaderMainFun)
{
	unsigned int ShaderCacheKey = _ComputeShaderCacheKey(pShaderFileName, pShaderMainFun);

	map<unsigned int, ID3D10Blob*>::iterator it = m_mapVSCodeBlob.find(ShaderCacheKey);

	if(it != m_mapVSCodeBlob.end())
		return it->second;
	else
		return 0;
}

ID3D10Blob* ShaderCompileManager::FindPSCodeBlob(const char* pShaderFileName, const char* pShaderMainFun)
{
	unsigned int ShaderCacheKey = _ComputeShaderCacheKey(pShaderFileName, pShaderMainFun);

	map<unsigned int, ID3D10Blob*>::iterator it = m_mapPSCodeBlob.find(ShaderCacheKey);

	if(it != m_mapPSCodeBlob.end())
		return it->second;
	else
		return 0;
}

unsigned int ShaderCompileManager::_ComputeShaderCacheKey(const char* pShaderFileName, const char* pShaderMainFun)
{
	char ShaderCacheString[SHADER_CACHE_SIZE] = {0};

	strcat(ShaderCacheString, pShaderFileName);

	strcat(ShaderCacheString, pShaderMainFun);

	return BKDRHash(ShaderCacheString);
}

ID3D11VertexShader* ShaderCompileManager::_FindVSShaderInCache(unsigned int ShaderCacheKey)
{
	map<unsigned int, ID3D11VertexShader*>::iterator it = m_mapVSCache.find(ShaderCacheKey);

	if(it != m_mapVSCache.end())
		return it->second;
	else
		return 0;
}

ID3D11PixelShader* ShaderCompileManager::_FindPSShaderInCache(unsigned int ShaderCacheKey)
{
	map<unsigned int, ID3D11PixelShader*>::iterator it = m_mapPSCache.find(ShaderCacheKey);

	if(it != m_mapPSCache.end())
		return it->second;
	else
		return 0;
}

string ShaderCompileManager::_ReadShaderToString(const char* pShaderFileName)
{
	char szShaderFilePath[MAX_PATH] = {0};
	
	GetCurrentDirectory(MAX_PATH, szShaderFilePath);
	
	strcat(szShaderFilePath, "\\EngineRes\\HLSLShader\\");

	strcat(szShaderFilePath, pShaderFileName);

	return ReadFileToString(szShaderFilePath);
}